package io.gitee.sfch.chunmi.common.tree;


import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import io.gitee.sfch.chunmi.common.tool.CollectionTool;
import lombok.Data;
import lombok.experimental.Accessors;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Supplier;

/**
 * {@link TreeNode} 建造器, 使用方法：
 * <blockquote><pre>
 *       // 菜单实体例子
 *       public class Menu {
 *          private String Id;
 *          private String parentId;
 *          private String menuName;
 *       }
 *       // 构建TreeBuilder
 *       List<BaseTreeNode<Menu>> treeNods = new TreeBuilder<Menu, BaseTreeNode<Menu>>(
 *          Menu::getId,
 *          Menu::getParentId,
 *          BaseTreeNode::new
 *       )
 *       .setNodeExtraSetter((node, data) -> {
 *          node.setTitle(data.getMenuName());
 *          node.setData(data);
 *       })
 *       .build(list);
 * </pre></blockquote><p>
 *
 * @param <T> 树节点数据类型
 * @param <N> 树节点类型，必须是 {@link TreeNode} 或其子类
 */
@Data
@Accessors(chain = true)
public class TreeBuilder<T, N extends TreeNode<N>> {

    /**
     * 获取根节点的父主键值（可选）, 默认根节点的父节点为空
     */
    private Object rootParentId = null;

    /**
     * 获取节点ID的方法
     */
    private Function<T, Object> idGetter;

    /**
     * 获取父节点ID的方法
     */
    private Function<T, Object> parentIdGetter;

    /**
     * 节点排序比较器（可选）
     */
    private Comparator<T> nodeComparator;

    /**
     * 为节点增加其他信息（可选）
     */
    private BiConsumer<N, T> nodeExtraSetter;

    /**
     * 创建节点实例的方法
     */
    private Supplier<N> nodeFactory;

    /**
     * 构建{@link TreeBuilder}
     *
     * @param idGetter 获取节点ID的方法，必填
     * @param parentIdGetter 获取父节点ID的方法，必填
     * @param nodeFactory 创建节点实例的方法，必填
     */
    public TreeBuilder(Function<T, Object> idGetter, Function<T, Object> parentIdGetter, Supplier<N> nodeFactory) {
        this.idGetter = idGetter;
        this.parentIdGetter = parentIdGetter;
        this.nodeFactory = nodeFactory;
    }

    /**
     * 构建节点
     *
     * @param nodeList 节点集合
     * @return 返回树的根节点
     */
    public List<N> build(List<T> nodeList) {
        // 没有节点数据的话直接返回空集合
        if (CollectionUtil.isEmpty(nodeList)) {
            return Collections.emptyList();
        }

        // 如果需要排序，对根节点进行排序
        if (nodeComparator != null) {
            nodeList.sort(nodeComparator);
        }

        // 按父节点主键进行分组
        Map<Object, List<T>> parentMap = CollectionTool.groupBy(nodeList, new LinkedHashMap<>(), getParentIdGetter());

        // 先从根节点开始构建
        Object rootParentId = this.rootParentId;
        // 基于remove可以有效的防御父ID循环引用
        List<T> roots = parentMap.remove(rootParentId);
        // 没有根节点的话返回空集合
        if (CollectionUtil.isEmpty(roots)) {
            return Collections.emptyList();
        }

        List<N> rootNodes = new ArrayList<>(roots.size());
        for (T rootData : roots) {
            N rootNode = buildTreeNode(rootData);
            rootNodes.add(rootNode);

            // 构建子节点
            buildChildren(rootNode, parentMap);
        }

        return rootNodes;
    }

    /**
     * 递归构建子节点
     *
     * @param parentNode 父节点
     * @param parentMap 父节点ID到子节点列表的映射
     */
    protected void buildChildren(N parentNode, Map<Object, List<T>> parentMap) {
        String parentId = parentNode.getId();
        // 基于remove可以有效的防御父ID循环引用
        List<T> children = parentMap.remove(parentId);

        if (CollectionUtil.isEmpty(children)) {
            return;
        }

        List<N> childNodes = new ArrayList<>(children.size());
        for (T childData : children) {
            N childNode = buildTreeNode(childData);
            childNodes.add(childNode);

            // 递归构建子节点的子节点
            buildChildren(childNode, parentMap);
        }

        parentNode.setChildren(childNodes);
    }

    /**
     * 构建单个树节点
     *
     * @param nodeData 节点数据
     * @return 构建好的树节点
     */
    protected N buildTreeNode(T nodeData) {
        N node = nodeFactory.get();
        node.setId(StrUtil.toStringOrNull(idGetter.apply(nodeData)));
        node.setParentId(StrUtil.toStringOrNull(parentIdGetter.apply(nodeData)));
        if (nodeExtraSetter != null) {
            nodeExtraSetter.accept(node, nodeData);
        }

        return node;
    }
}

